diff --git a/cpp/include/cudf/io/config_utils.hpp b/cpp/include/cudf/io/config_utils.hpp index 20a6747f6ca..38001c71a76 100644 --- a/cpp/include/cudf/io/config_utils.hpp +++ b/cpp/include/cudf/io/config_utils.hpp @@ -6,6 +6,9 @@ #include +#include +#include + /** * @file * @brief APIs for configuring KvikIO-based I/O behavior. @@ -27,10 +30,12 @@ namespace kvikio_integration { * Parameters include: * - Compatibility mode, according to the environment variable KVIKIO_COMPAT_MODE. If * KVIKIO_COMPAT_MODE is not set, enable it by default, which enforces the use of POSIX I/O. - * - Thread pool size, according to the environment variable KVIKIO_NTHREADS. If KVIKIO_NTHREADS is - * not set, use 4 threads by default. + * - Thread pool size. If @p nthreads is provided, it is used directly. Otherwise, the value is + * read from the environment variable KVIKIO_NTHREADS, defaulting to 4 if unset. + * + * @param nthreads Optional thread pool size override. If provided, supersedes KVIKIO_NTHREADS. */ -void set_up_kvikio(); +void set_up_kvikio(std::optional nthreads = std::nullopt); } // namespace kvikio_integration diff --git a/cpp/src/io/utilities/config_utils.cpp b/cpp/src/io/utilities/config_utils.cpp index a7ba392cf72..7f6511f06ed 100644 --- a/cpp/src/io/utilities/config_utils.cpp +++ b/cpp/src/io/utilities/config_utils.cpp @@ -10,16 +10,18 @@ #include +#include +#include #include namespace cudf::io { namespace kvikio_integration { -void set_up_kvikio() +void set_up_kvikio(std::optional nthreads) { static std::once_flag flag{}; - std::call_once(flag, [] { + std::call_once(flag, [nthreads] { // Workaround for https://github.com/NVIDIA/cudf/issues/14140, where cuFileDriverOpen errors // out if no CUDA calls have been made before it. This is a no-op if the CUDA context is already // initialized. @@ -28,8 +30,8 @@ void set_up_kvikio() auto const compat_mode = kvikio::getenv_or("KVIKIO_COMPAT_MODE", kvikio::CompatMode::ON); kvikio::defaults::set_compat_mode(compat_mode); - auto const nthreads = cudf::detail::getenv_or("KVIKIO_NTHREADS", 4u); - kvikio::defaults::set_thread_pool_nthreads(nthreads); + auto const n = nthreads.value_or(cudf::detail::getenv_or("KVIKIO_NTHREADS", 4u)); + kvikio::defaults::set_thread_pool_nthreads(n); }); } diff --git a/python/cudf_polars/cudf_polars/engine/core.py b/python/cudf_polars/cudf_polars/engine/core.py index 07fe1750d64..f73e11bdca3 100644 --- a/python/cudf_polars/cudf_polars/engine/core.py +++ b/python/cudf_polars/cudf_polars/engine/core.py @@ -286,12 +286,12 @@ class StreamingEngine(pl.GPUEngine): destruction and context manager exit must occur on the thread that created the instance. - Creating an engine configures the process-wide kvikio thread pool (default - 256 threads). Because kvikio's pool is a global singleton, this blocks - any concurrent kvikio IO in the process until in-flight IO completes and overrides any prior - ``kvikio.defaults.set("num_threads", ...)`` call. Use the - ``kvikio_nthreads`` executor option or the ``KVIKIO_NTHREADS`` environment - variable to control the thread count. + Creating an engine sets the kvikio remote I/O backend to ``EASY_THREADPOOL`` + and configures its thread pool (default 256 threads). Because kvikio's pool + is a global singleton, this blocks any concurrent kvikio IO in the process + until in-flight IO completes and overrides any prior ``kvikio.defaults.set(...)`` + calls. Use the ``kvikio_nthreads`` executor option or the ``KVIKIO_NTHREADS`` + environment variable to control the thread count. Parameters ---------- diff --git a/python/cudf_polars/cudf_polars/engine/dask.py b/python/cudf_polars/cudf_polars/engine/dask.py index 7e3d1fa0c20..a2154a78407 100644 --- a/python/cudf_polars/cudf_polars/engine/dask.py +++ b/python/cudf_polars/cudf_polars/engine/dask.py @@ -15,8 +15,6 @@ import distributed import distributed.system -import kvikio -import kvikio.defaults import pynvml import ucxx._lib.libucxx as ucx_api @@ -59,6 +57,7 @@ from cudf_polars.utils.config import ( DaskContext, MemoryResourceConfig, + configure_kvikio, resolve_kvikio_nthreads, resolve_kvikio_statistics, ) @@ -66,6 +65,8 @@ if TYPE_CHECKING: from collections.abc import Callable + import kvikio + from cudf_streaming.channel_metadata import ChannelMetadata from rapidsmpf.communicator.communicator import Communicator from rapidsmpf.rmm_resource_adaptor import RmmResourceAdaptor @@ -383,7 +384,6 @@ def _setup_worker( """ assert dask_worker is not None - kvikio.defaults.set("num_threads", kvikio_nthreads) options = Options.deserialize(rapidsmpf_options_as_bytes) attr = f"_cudf_polars_mp_context_{uid}" mp_ctx: _WorkerContext | None = getattr(dask_worker, attr, None) @@ -391,6 +391,7 @@ def _setup_worker( if mp_ctx is None: # Non-root worker: create communicator now. bind_to_gpu(hardware_binding) + configure_kvikio(kvikio_nthreads) memory_resource_config = ( memory_resource_config or MemoryResourceConfig.default() ) @@ -411,6 +412,7 @@ def _setup_worker( base_mr = mp_ctx.base_mr comm = mp_ctx.comm statistics = mp_ctx.statistics + configure_kvikio(kvikio_nthreads) barrier(comm) worker_id = worker_ids[comm.rank] @@ -530,7 +532,7 @@ def _reset_worker( Injected by ``distributed`` when called via :meth:`distributed.Client.run`. """ assert dask_worker is not None - kvikio.defaults.set("num_threads", kvikio_nthreads) + configure_kvikio(kvikio_nthreads) attr = f"_cudf_polars_mp_context_{uid}" mp_ctx: _WorkerContext | None = getattr(dask_worker, attr, None) if mp_ctx is None: diff --git a/python/cudf_polars/cudf_polars/engine/ray.py b/python/cudf_polars/cudf_polars/engine/ray.py index 42448399960..9a03e6e5f8b 100644 --- a/python/cudf_polars/cudf_polars/engine/ray.py +++ b/python/cudf_polars/cudf_polars/engine/ray.py @@ -10,8 +10,6 @@ from concurrent.futures import ThreadPoolExecutor from typing import TYPE_CHECKING, Any, cast -import kvikio -import kvikio.defaults import ray import ray.exceptions import ucxx._lib.libucxx as ucx_api @@ -56,6 +54,7 @@ from cudf_polars.utils.config import ( MemoryResourceConfig, RayContext, + configure_kvikio, resolve_kvikio_nthreads, resolve_kvikio_statistics, ) @@ -63,6 +62,7 @@ if TYPE_CHECKING: from collections.abc import Callable + import kvikio from ray import ObjectRef from ray.actor import ActorHandle @@ -261,7 +261,7 @@ def __init__( quent_enabled: bool, ) -> None: bind_to_gpu(hardware_binding) - kvikio.defaults.set("num_threads", kvikio_nthreads) + configure_kvikio(kvikio_nthreads) memory_resource_config = ( memory_resource_config or MemoryResourceConfig.default() ) @@ -383,7 +383,7 @@ def reset( """ if self._ctx is None: raise RuntimeError("reset() requires setup_worker() to have run") - kvikio.defaults.set("num_threads", kvikio_nthreads) + configure_kvikio(kvikio_nthreads) assert self._comm is not None # Collective: all ranks idle before any rank tears down its Context. if self._comm.nranks > 1: diff --git a/python/cudf_polars/cudf_polars/engine/spmd.py b/python/cudf_polars/cudf_polars/engine/spmd.py index 38d712f9bbb..6e2ab5ad0c8 100644 --- a/python/cudf_polars/cudf_polars/engine/spmd.py +++ b/python/cudf_polars/cudf_polars/engine/spmd.py @@ -64,6 +64,7 @@ MemoryResourceConfig, SPMDContext, StreamingExecutor, + configure_kvikio, resolve_kvikio_nthreads, resolve_kvikio_statistics, ) @@ -438,7 +439,7 @@ def __init__( ) bind_to_gpu(hw_binding) - kvikio.defaults.set("num_threads", executor_options["kvikio_nthreads"]) + configure_kvikio(executor_options["kvikio_nthreads"]) self.rapidsmpf_options = resolve_rapidsmpf_options(rapidsmpf_options) mr_config: MemoryResourceConfig = engine_options.get( @@ -629,10 +630,10 @@ def _reset( existing_kvikio_nthreads = existing_executor_options.get("kvikio_nthreads") if existing_kvikio_nthreads is not None: executor_options.setdefault("kvikio_nthreads", existing_kvikio_nthreads) + configure_kvikio(executor_options["kvikio_nthreads"]) executor_options.setdefault( "kvikio_statistics", resolve_kvikio_statistics(executor_options) ) - kvikio.defaults.set("num_threads", executor_options["kvikio_nthreads"]) engine_options = engine_options or {} quent_context: cudf_polars.quent.QuentContext | None = executor_options.get( "quent_context" diff --git a/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds.py b/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds.py index 19ac159b76e..5f921493493 100644 --- a/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds.py +++ b/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds.py @@ -37,7 +37,6 @@ # Without this setting, the first IO task to run # on each worker takes ~15 sec extra os.environ["KVIKIO_COMPAT_MODE"] = os.environ.get("KVIKIO_COMPAT_MODE", "on") -os.environ["KVIKIO_NTHREADS"] = os.environ.get("KVIKIO_NTHREADS", "8") # TODO: consider raising the rapidsmpf built-in default from 1 to 8. os.environ["RAPIDSMPF_NUM_STREAMING_THREADS"] = os.environ.get( "RAPIDSMPF_NUM_STREAMING_THREADS", "8" diff --git a/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsh.py b/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsh.py index f51d05d6c23..c4c8c28ce44 100644 --- a/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsh.py +++ b/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsh.py @@ -42,7 +42,6 @@ # Without this setting, the first IO task to run # on each worker takes ~15 sec extra os.environ["KVIKIO_COMPAT_MODE"] = os.environ.get("KVIKIO_COMPAT_MODE", "on") -os.environ["KVIKIO_NTHREADS"] = os.environ.get("KVIKIO_NTHREADS", "8") # TODO: consider raising the rapidsmpf built-in default from 1 to 8. os.environ["RAPIDSMPF_NUM_STREAMING_THREADS"] = os.environ.get( "RAPIDSMPF_NUM_STREAMING_THREADS", "8" diff --git a/python/cudf_polars/cudf_polars/utils/config.py b/python/cudf_polars/cudf_polars/utils/config.py index 38d9bef261d..6817ec5c162 100644 --- a/python/cudf_polars/cudf_polars/utils/config.py +++ b/python/cudf_polars/cudf_polars/utils/config.py @@ -29,6 +29,11 @@ import os from typing import TYPE_CHECKING, Any, Generic, Literal, TypeVar +import kvikio +import kvikio.defaults + +import pylibcudf.utils + if TYPE_CHECKING: import uuid from collections.abc import Callable @@ -221,6 +226,22 @@ def resolve_kvikio_nthreads(executor_options: dict[str, Any]) -> int: ) +def configure_kvikio(nthreads: int) -> None: + """Set the remote I/O backend to ``EASY_THREADPOOL`` with ``nthreads`` threads.""" + # HACK: libcudf calls set_up_kvikio() on the first IO op and that resets the thread + # pool (default is 4 if KVIKIO_NTHREADS is unset), undoing anything we set via + # kvikio.defaults. We call it here with our nthreads so later when it's called in + # libcudf it's a no-op. The explicit kvikio.defaults.set below handles subsequent + # calls to configure_kvikio (call_once only fires once). + pylibcudf.utils._set_up_kvikio(nthreads) + kvikio.defaults.set( + { + "num_threads": nthreads, + "remote_io_backend": kvikio.RemoteIOBackend.EASY_THREADPOOL, + } + ) + + def _bool_converter(v: str) -> bool: lowered = v.lower() if lowered in {"true", "yes", "y", "1"}: @@ -749,8 +770,9 @@ class StreamingExecutor: Maximum number of workers for the Python ThreadPoolExecutor. Default is 8. kvikio_nthreads - Number of threads in the kvikio thread pool. Defaults to 256, which is - tuned for cloud object-store IO. This can be set via + Number of threads in the kvikio ``EASY_THREADPOOL`` thread pool. + Defaults to 256, which is tuned for cloud object-store IO. This can be + set via - ``executor_options`` passed to ``polars.GPUEngine`` - the ``CUDF_POLARS__EXECUTOR__KVIKIO_NTHREADS`` environment variable diff --git a/python/cudf_polars/tests/test_config.py b/python/cudf_polars/tests/test_config.py index fd184cf2f2c..a0d2b05d1f6 100644 --- a/python/cudf_polars/tests/test_config.py +++ b/python/cudf_polars/tests/test_config.py @@ -37,6 +37,7 @@ ParquetOptions, StreamingExecutor, Unspecified, + configure_kvikio, ) from cudf_polars.utils.cuda_stream import get_cuda_stream @@ -920,6 +921,21 @@ def test_kvikio_nthreads_cudf_polars_env_takes_precedence( assert config.executor.kvikio_nthreads == 64 +def test_configure_kvikio_sets_backend_and_threads( + monkeypatch: pytest.MonkeyPatch, +) -> None: + import kvikio + import kvikio.defaults + + monkeypatch.delenv("KVIKIO_NTHREADS", raising=False) + configure_kvikio(42) + assert kvikio.defaults.get("num_threads") == 42 + assert ( + kvikio.defaults.get("remote_io_backend") + == kvikio.RemoteIOBackend.EASY_THREADPOOL + ) + + def test_dask_sink_to_directory_false_raises() -> None: with pytest.raises( ValueError, match="The dask cluster requires sink_to_directory=True" diff --git a/python/pylibcudf/pylibcudf/io/__init__.pxd b/python/pylibcudf/pylibcudf/io/__init__.pxd index d8a3c42d4c1..be38ba9fd2e 100644 --- a/python/pylibcudf/pylibcudf/io/__init__.pxd +++ b/python/pylibcudf/pylibcudf/io/__init__.pxd @@ -1,4 +1,4 @@ -# SPDX-FileCopyrightText: Copyright (c) 2024-2025, NVIDIA CORPORATION. +# SPDX-FileCopyrightText: Copyright (c) 2024-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 # CSV is removed since it is def not cpdef (to force kw-only arguments) diff --git a/python/pylibcudf/pylibcudf/io/__init__.py b/python/pylibcudf/pylibcudf/io/__init__.py index 1f0a0a21819..1efe8939764 100644 --- a/python/pylibcudf/pylibcudf/io/__init__.py +++ b/python/pylibcudf/pylibcudf/io/__init__.py @@ -29,6 +29,7 @@ "datasource", "experimental", "json", + "kvikio", "orc", "parquet", "parquet_io_utils", diff --git a/python/pylibcudf/pylibcudf/libcudf/utilities/config_utils.pxd b/python/pylibcudf/pylibcudf/libcudf/utilities/config_utils.pxd new file mode 100644 index 00000000000..af7babb2382 --- /dev/null +++ b/python/pylibcudf/pylibcudf/libcudf/utilities/config_utils.pxd @@ -0,0 +1,11 @@ +# SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from libc.stdint cimport uint32_t +from libcpp.optional cimport optional +from pylibcudf.exception_handler cimport libcudf_exception_handler + +cdef extern from "cudf/io/config_utils.hpp" \ + namespace "cudf::io::kvikio_integration" nogil: + + void set_up_kvikio(optional[uint32_t] nthreads) except +libcudf_exception_handler diff --git a/python/pylibcudf/pylibcudf/utils.pxd b/python/pylibcudf/pylibcudf/utils.pxd index feb82cea18f..c1d49d290af 100644 --- a/python/pylibcudf/pylibcudf/utils.pxd +++ b/python/pylibcudf/pylibcudf/utils.pxd @@ -1,4 +1,4 @@ -# SPDX-FileCopyrightText: Copyright (c) 2023-2026, NVIDIA CORPORATION. +# SPDX-FileCopyrightText: Copyright (c) 2023-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 from libcpp.functional cimport reference_wrapper @@ -10,3 +10,4 @@ from rmm.pylibrmm.stream cimport Stream cdef vector[reference_wrapper[const scalar]] _as_vector(list source) cpdef Stream _get_stream(object stream = *) cdef DeviceMemoryResource _get_memory_resource(DeviceMemoryResource mr = *) +cpdef void _set_up_kvikio(object nthreads = *) diff --git a/python/pylibcudf/pylibcudf/utils.pyi b/python/pylibcudf/pylibcudf/utils.pyi index bc3898629e2..d10f5524b6d 100644 --- a/python/pylibcudf/pylibcudf/utils.pyi +++ b/python/pylibcudf/pylibcudf/utils.pyi @@ -9,3 +9,5 @@ class HasCudaStream(Protocol): def __cuda_stream__(self) -> tuple[int, int]: ... CudaStreamLike = Stream | HasCudaStream + +def _set_up_kvikio(nthreads: int | None = None) -> None: ... diff --git a/python/pylibcudf/pylibcudf/utils.pyx b/python/pylibcudf/pylibcudf/utils.pyx index 0bae89fd19a..f38cba03c71 100644 --- a/python/pylibcudf/pylibcudf/utils.pyx +++ b/python/pylibcudf/pylibcudf/utils.pyx @@ -3,8 +3,11 @@ from cython.operator import dereference +from libc.stdint cimport uint32_t from libcpp.functional cimport reference_wrapper +from libcpp.optional cimport make_optional, nullopt, optional from libcpp.vector cimport vector +from pylibcudf.libcudf.utilities.config_utils cimport set_up_kvikio as cpp_set_up_kvikio from pylibcudf.libcudf.scalar.scalar cimport scalar @@ -67,3 +70,11 @@ cdef DeviceMemoryResource _get_memory_resource(DeviceMemoryResource mr = None): if mr is None: return get_current_device_resource() return mr + + +cpdef void _set_up_kvikio(object nthreads=None): + cdef optional[uint32_t] c_nthreads = nullopt + if nthreads is not None: + c_nthreads = make_optional[uint32_t](nthreads) + with nogil: + cpp_set_up_kvikio(c_nthreads)