diff --git a/cpp/include/cudf/io/config_utils.hpp b/cpp/include/cudf/io/config_utils.hpp index 20a6747f6ca0..38001c71a76b 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 a7ba392cf724..7f6511f06ed5 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/docs/cudf/source/pylibcudf/api_docs/io/index.rst b/docs/cudf/source/pylibcudf/api_docs/io/index.rst index 45b4def70051..521cda7d0ab7 100644 --- a/docs/cudf/source/pylibcudf/api_docs/io/index.rst +++ b/docs/cudf/source/pylibcudf/api_docs/io/index.rst @@ -19,6 +19,7 @@ I/O Functions csv experimental json + kvikio orc parquet parquet_io_utils diff --git a/docs/cudf/source/pylibcudf/api_docs/io/kvikio.rst b/docs/cudf/source/pylibcudf/api_docs/io/kvikio.rst new file mode 100644 index 000000000000..b06b030f19a1 --- /dev/null +++ b/docs/cudf/source/pylibcudf/api_docs/io/kvikio.rst @@ -0,0 +1,6 @@ +====== +KvikIO +====== + +.. automodule:: pylibcudf.io.kvikio + :members: diff --git a/python/cudf_polars/cudf_polars/engine/core.py b/python/cudf_polars/cudf_polars/engine/core.py index 3a58ef3ac1c3..a9e482b568b3 100644 --- a/python/cudf_polars/cudf_polars/engine/core.py +++ b/python/cudf_polars/cudf_polars/engine/core.py @@ -185,12 +185,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 86a409a7413e..916ad2cafb25 100644 --- a/python/cudf_polars/cudf_polars/engine/dask.py +++ b/python/cudf_polars/cudf_polars/engine/dask.py @@ -15,7 +15,6 @@ import distributed import distributed.system -import kvikio.defaults import pynvml import ucxx._lib.libucxx as ucx_api @@ -55,6 +54,7 @@ from cudf_polars.utils.config import ( DaskContext, MemoryResourceConfig, + configure_kvikio, resolve_kvikio_nthreads, ) @@ -374,7 +374,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) @@ -382,6 +381,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() ) @@ -402,6 +402,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] @@ -512,7 +513,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 38b088756926..c7ff4943f40a 100644 --- a/python/cudf_polars/cudf_polars/engine/ray.py +++ b/python/cudf_polars/cudf_polars/engine/ray.py @@ -10,7 +10,6 @@ from concurrent.futures import ThreadPoolExecutor from typing import TYPE_CHECKING, Any, cast -import kvikio.defaults import ray import ray.exceptions import ucxx._lib.libucxx as ucx_api @@ -52,6 +51,7 @@ from cudf_polars.utils.config import ( MemoryResourceConfig, RayContext, + configure_kvikio, resolve_kvikio_nthreads, ) @@ -255,7 +255,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() ) @@ -368,7 +368,7 @@ def reset(self, *, rapidsmpf_options_as_bytes: bytes, kvikio_nthreads: int) -> N """ 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 10db61bd388c..6468d42c4ecb 100644 --- a/python/cudf_polars/cudf_polars/engine/spmd.py +++ b/python/cudf_polars/cudf_polars/engine/spmd.py @@ -11,8 +11,6 @@ from concurrent.futures import ThreadPoolExecutor from typing import TYPE_CHECKING, Any, cast -import kvikio.defaults - import pylibcudf as plc import rmm.mr from cudf_streaming.partition_utils import ( @@ -60,6 +58,7 @@ MemoryResourceConfig, SPMDContext, StreamingExecutor, + configure_kvikio, resolve_kvikio_nthreads, ) @@ -430,7 +429,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( @@ -611,7 +610,7 @@ 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) - kvikio.defaults.set("num_threads", executor_options["kvikio_nthreads"]) + configure_kvikio(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 19ac159b76e0..5f921493493c 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 f51d05d6c23f..c4c8c28ce441 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 315aa0c8f8cd..cd7a03f22db9 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.io.kvikio + if TYPE_CHECKING: import uuid from collections.abc import Callable @@ -211,6 +216,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.io.kvikio.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"}: @@ -739,8 +760,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 fd184cf2f2c3..a0d2b05d1f62 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/CMakeLists.txt b/python/pylibcudf/pylibcudf/io/CMakeLists.txt index 65bb31d908ff..b5984c8b013c 100644 --- a/python/pylibcudf/pylibcudf/io/CMakeLists.txt +++ b/python/pylibcudf/pylibcudf/io/CMakeLists.txt @@ -5,7 +5,7 @@ # cmake-format: on # ============================================================================= -set(cython_sources avro.pyx csv.pyx datasource.pyx json.pyx orc.pyx parquet.pyx +set(cython_sources avro.pyx csv.pyx datasource.pyx json.pyx kvikio.pyx orc.pyx parquet.pyx parquet_io_utils.pyx parquet_metadata.pyx text.pyx timezone.pyx types.pyx ) diff --git a/python/pylibcudf/pylibcudf/io/__init__.pxd b/python/pylibcudf/pylibcudf/io/__init__.pxd index d8a3c42d4c1b..3e8c7984cc28 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-2026, 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) @@ -7,6 +7,7 @@ from . cimport ( datasource, experimental, json, + kvikio, orc, parquet, parquet_metadata, diff --git a/python/pylibcudf/pylibcudf/io/__init__.py b/python/pylibcudf/pylibcudf/io/__init__.py index 1f0a0a218199..607656222155 100644 --- a/python/pylibcudf/pylibcudf/io/__init__.py +++ b/python/pylibcudf/pylibcudf/io/__init__.py @@ -7,6 +7,7 @@ datasource, experimental, json, + kvikio, orc, parquet, parquet_io_utils, @@ -29,6 +30,7 @@ "datasource", "experimental", "json", + "kvikio", "orc", "parquet", "parquet_io_utils", diff --git a/python/pylibcudf/pylibcudf/io/kvikio.pxd b/python/pylibcudf/pylibcudf/io/kvikio.pxd new file mode 100644 index 000000000000..7a2b48d8d551 --- /dev/null +++ b/python/pylibcudf/pylibcudf/io/kvikio.pxd @@ -0,0 +1,4 @@ +# SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +cpdef void set_up_kvikio(object nthreads=*) diff --git a/python/pylibcudf/pylibcudf/io/kvikio.pyi b/python/pylibcudf/pylibcudf/io/kvikio.pyi new file mode 100644 index 000000000000..4903454603eb --- /dev/null +++ b/python/pylibcudf/pylibcudf/io/kvikio.pyi @@ -0,0 +1,4 @@ +# SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +def set_up_kvikio(nthreads: int | None = None) -> None: ... diff --git a/python/pylibcudf/pylibcudf/io/kvikio.pyx b/python/pylibcudf/pylibcudf/io/kvikio.pyx new file mode 100644 index 000000000000..b797a9bab2e1 --- /dev/null +++ b/python/pylibcudf/pylibcudf/io/kvikio.pyx @@ -0,0 +1,34 @@ +# 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 make_optional, nullopt, optional +from pylibcudf.libcudf.io.config_utils cimport set_up_kvikio as cpp_set_up_kvikio + +__all__ = ["set_up_kvikio"] + + +cpdef void set_up_kvikio(object nthreads=None): + """Set KvikIO parameters. + + 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. If ``nthreads`` is provided, it is used directly. Otherwise, the value is + read from the environment variable ``KVIKIO_NTHREADS``, defaulting to 4 if unset. + + Parameters + ---------- + nthreads : int, optional + Thread pool size override. If provided, supersedes ``KVIKIO_NTHREADS``. + + Returns + ------- + 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) diff --git a/python/pylibcudf/pylibcudf/libcudf/io/config_utils.pxd b/python/pylibcudf/pylibcudf/libcudf/io/config_utils.pxd new file mode 100644 index 000000000000..af7babb23820 --- /dev/null +++ b/python/pylibcudf/pylibcudf/libcudf/io/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