From a541d88e63d9bfc4d2853c704207451d5e040b43 Mon Sep 17 00:00:00 2001 From: Matthew Murray Date: Mon, 17 Aug 2026 18:23:13 +0000 Subject: [PATCH 1/5] Set default kvikio backend in cudf-polars to EASY_THREADPOOL --- python/cudf_polars/cudf_polars/engine/core.py | 11 ++++++----- python/cudf_polars/cudf_polars/engine/dask.py | 6 +++--- python/cudf_polars/cudf_polars/engine/ray.py | 6 +++--- python/cudf_polars/cudf_polars/engine/spmd.py | 7 +++---- python/cudf_polars/cudf_polars/utils/config.py | 13 +++++++++++++ python/cudf_polars/tests/test_config.py | 13 +++++++++++++ 6 files changed, 41 insertions(+), 15 deletions(-) diff --git a/python/cudf_polars/cudf_polars/engine/core.py b/python/cudf_polars/cudf_polars/engine/core.py index 3a58ef3ac1c3..ac7ff851ab5f 100644 --- a/python/cudf_polars/cudf_polars/engine/core.py +++ b/python/cudf_polars/cudf_polars/engine/core.py @@ -186,11 +186,12 @@ class StreamingEngine(pl.GPUEngine): 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. + 256 threads) and sets the remote I/O backend to ``EASY_THREADPOOL``. 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 649046997880..7552a1f63239 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,7 @@ def _setup_worker( """ assert dask_worker is not None - kvikio.defaults.set("num_threads", kvikio_nthreads) + configure_kvikio(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) @@ -512,7 +512,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 0b4d4078f821..15561aed34aa 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, ) @@ -252,7 +252,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() ) @@ -365,7 +365,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/utils/config.py b/python/cudf_polars/cudf_polars/utils/config.py index b7aeabf58140..848fedecea22 100644 --- a/python/cudf_polars/cudf_polars/utils/config.py +++ b/python/cudf_polars/cudf_polars/utils/config.py @@ -29,6 +29,9 @@ import os from typing import TYPE_CHECKING, Any, Generic, Literal, TypeVar +import kvikio +import kvikio.defaults + if TYPE_CHECKING: import uuid from collections.abc import Callable @@ -211,6 +214,16 @@ def resolve_kvikio_nthreads(executor_options: dict[str, Any]) -> int: ) +def configure_kvikio(nthreads: int) -> None: + """Configure the process-wide kvikio thread pool and remote I/O backend.""" + 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"}: diff --git a/python/cudf_polars/tests/test_config.py b/python/cudf_polars/tests/test_config.py index 59321b07cf39..50d1d765bfc9 100644 --- a/python/cudf_polars/tests/test_config.py +++ b/python/cudf_polars/tests/test_config.py @@ -36,6 +36,7 @@ ParquetOptions, StreamingExecutor, Unspecified, + configure_kvikio, ) from cudf_polars.utils.cuda_stream import get_cuda_stream @@ -909,6 +910,18 @@ def test_kvikio_nthreads_cudf_polars_env_takes_precedence( assert config.executor.kvikio_nthreads == 64 +def test_configure_kvikio_sets_backend_and_threads() -> None: + import kvikio + import kvikio.defaults + + 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" From 75b924bc5b182dc7a7b7f5c2a1dd215310368cc1 Mon Sep 17 00:00:00 2001 From: Matthew Murray Date: Mon, 17 Aug 2026 18:23:41 +0000 Subject: [PATCH 2/5] docs --- python/cudf_polars/cudf_polars/engine/core.py | 13 ++++++------- python/cudf_polars/cudf_polars/utils/config.py | 7 ++++--- 2 files changed, 10 insertions(+), 10 deletions(-) diff --git a/python/cudf_polars/cudf_polars/engine/core.py b/python/cudf_polars/cudf_polars/engine/core.py index ac7ff851ab5f..a9e482b568b3 100644 --- a/python/cudf_polars/cudf_polars/engine/core.py +++ b/python/cudf_polars/cudf_polars/engine/core.py @@ -185,13 +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) and sets the remote I/O backend to ``EASY_THREADPOOL``. 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. + 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/utils/config.py b/python/cudf_polars/cudf_polars/utils/config.py index 848fedecea22..176a3471a2ec 100644 --- a/python/cudf_polars/cudf_polars/utils/config.py +++ b/python/cudf_polars/cudf_polars/utils/config.py @@ -215,7 +215,7 @@ def resolve_kvikio_nthreads(executor_options: dict[str, Any]) -> int: def configure_kvikio(nthreads: int) -> None: - """Configure the process-wide kvikio thread pool and remote I/O backend.""" + """Set the remote I/O backend to ``EASY_THREADPOOL`` with ``nthreads`` threads.""" kvikio.defaults.set( { "num_threads": nthreads, @@ -752,8 +752,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 From 344e83de6435ffded36fe378eac858f3d091f578 Mon Sep 17 00:00:00 2001 From: Matthew Murray Date: Mon, 17 Aug 2026 21:09:49 +0000 Subject: [PATCH 3/5] fix interation with libcudfs setting of kvikio thread pool --- .../source/pylibcudf/api_docs/io/index.rst | 1 + .../source/pylibcudf/api_docs/io/kvikio.rst | 6 +++++ .../cudf_polars/streaming/benchmarks/pdsds.py | 1 - .../cudf_polars/streaming/benchmarks/pdsh.py | 1 - .../cudf_polars/cudf_polars/utils/config.py | 9 +++++++ python/cudf_polars/tests/test_config.py | 5 +++- python/pylibcudf/pylibcudf/io/CMakeLists.txt | 4 ++-- python/pylibcudf/pylibcudf/io/__init__.pxd | 3 ++- python/pylibcudf/pylibcudf/io/__init__.py | 2 ++ python/pylibcudf/pylibcudf/io/kvikio.pxd | 4 ++++ python/pylibcudf/pylibcudf/io/kvikio.pyi | 4 ++++ python/pylibcudf/pylibcudf/io/kvikio.pyx | 24 +++++++++++++++++++ .../pylibcudf/libcudf/io/config_utils.pxd | 9 +++++++ 13 files changed, 67 insertions(+), 6 deletions(-) create mode 100644 docs/cudf/source/pylibcudf/api_docs/io/kvikio.rst create mode 100644 python/pylibcudf/pylibcudf/io/kvikio.pxd create mode 100644 python/pylibcudf/pylibcudf/io/kvikio.pyi create mode 100644 python/pylibcudf/pylibcudf/io/kvikio.pyx create mode 100644 python/pylibcudf/pylibcudf/libcudf/io/config_utils.pxd diff --git a/docs/cudf/source/pylibcudf/api_docs/io/index.rst b/docs/cudf/source/pylibcudf/api_docs/io/index.rst index 15a87175325f..024735eceb61 100644 --- a/docs/cudf/source/pylibcudf/api_docs/io/index.rst +++ b/docs/cudf/source/pylibcudf/api_docs/io/index.rst @@ -18,6 +18,7 @@ I/O Functions avro csv json + kvikio orc parquet parquet_metadata 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/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 176a3471a2ec..7c58f2052365 100644 --- a/python/cudf_polars/cudf_polars/utils/config.py +++ b/python/cudf_polars/cudf_polars/utils/config.py @@ -32,6 +32,8 @@ import kvikio import kvikio.defaults +import pylibcudf.io.kvikio + if TYPE_CHECKING: import uuid from collections.abc import Callable @@ -216,6 +218,13 @@ 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: cudf calls set_up_kvikio() lazily on the first IO operation via std::call_once. + # That call reads KVIKIO_NTHREADS from the environment and resets the thread pool, + # undoing anything we set via kvikio.defaults. Setting the env var here means + # set_up_kvikio() will size the pool correctly when it fires, and calling it eagerly + # satisfies the once-flag so it becomes a no-op on the first IO. + os.environ["KVIKIO_NTHREADS"] = str(nthreads) + pylibcudf.io.kvikio.set_up_kvikio() kvikio.defaults.set( { "num_threads": nthreads, diff --git a/python/cudf_polars/tests/test_config.py b/python/cudf_polars/tests/test_config.py index 50d1d765bfc9..de6cb7d0edf2 100644 --- a/python/cudf_polars/tests/test_config.py +++ b/python/cudf_polars/tests/test_config.py @@ -910,10 +910,13 @@ def test_kvikio_nthreads_cudf_polars_env_takes_precedence( assert config.executor.kvikio_nthreads == 64 -def test_configure_kvikio_sets_backend_and_threads() -> None: +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 ( diff --git a/python/pylibcudf/pylibcudf/io/CMakeLists.txt b/python/pylibcudf/pylibcudf/io/CMakeLists.txt index 089ea8d0e8d9..0b611f2e1477 100644 --- a/python/pylibcudf/pylibcudf/io/CMakeLists.txt +++ b/python/pylibcudf/pylibcudf/io/CMakeLists.txt @@ -1,11 +1,11 @@ # ============================================================================= # cmake-format: off -# 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 # 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_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 a6a0ebad3a1e..2d222ee8c484 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_metadata, @@ -28,6 +29,7 @@ "datasource", "experimental", "json", + "kvikio", "orc", "parquet", "parquet_metadata", diff --git a/python/pylibcudf/pylibcudf/io/kvikio.pxd b/python/pylibcudf/pylibcudf/io/kvikio.pxd new file mode 100644 index 000000000000..6f1e861bc6ba --- /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() diff --git a/python/pylibcudf/pylibcudf/io/kvikio.pyi b/python/pylibcudf/pylibcudf/io/kvikio.pyi new file mode 100644 index 000000000000..3ff97c90c9de --- /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() -> None: ... diff --git a/python/pylibcudf/pylibcudf/io/kvikio.pyx b/python/pylibcudf/pylibcudf/io/kvikio.pyx new file mode 100644 index 000000000000..2a559344680c --- /dev/null +++ b/python/pylibcudf/pylibcudf/io/kvikio.pyx @@ -0,0 +1,24 @@ +# SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from pylibcudf.libcudf.io.config_utils cimport set_up_kvikio as cpp_set_up_kvikio + +__all__ = ["set_up_kvikio"] + + +cpdef void set_up_kvikio(): + """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, according to the environment variable ``KVIKIO_NTHREADS``. If + ``KVIKIO_NTHREADS`` is not set, use 4 threads by default. + + Returns + ------- + None + """ + with nogil: + cpp_set_up_kvikio() 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..c6d05b21ca31 --- /dev/null +++ b/python/pylibcudf/pylibcudf/libcudf/io/config_utils.pxd @@ -0,0 +1,9 @@ +# SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +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() except +libcudf_exception_handler From ee3bf0aa2ed11944abcd6c8dbdb74956cab7cd24 Mon Sep 17 00:00:00 2001 From: Matthew Murray Date: Mon, 17 Aug 2026 21:40:37 +0000 Subject: [PATCH 4/5] hack --- python/cudf_polars/cudf_polars/utils/config.py | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/python/cudf_polars/cudf_polars/utils/config.py b/python/cudf_polars/cudf_polars/utils/config.py index 7c58f2052365..a0fc72dcf0f6 100644 --- a/python/cudf_polars/cudf_polars/utils/config.py +++ b/python/cudf_polars/cudf_polars/utils/config.py @@ -218,11 +218,11 @@ 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: cudf calls set_up_kvikio() lazily on the first IO operation via std::call_once. - # That call reads KVIKIO_NTHREADS from the environment and resets the thread pool, - # undoing anything we set via kvikio.defaults. Setting the env var here means - # set_up_kvikio() will size the pool correctly when it fires, and calling it eagerly - # satisfies the once-flag so it becomes a no-op on the first IO. + # HACK: libcudf calls set_up_kvikio() on the first IO op and that resets the thread + # pool by reading KVIKIO_NTHREADS (default is 4 if unset), undoing anything we set + # via kvikio.defaults. We set KVIKIO_NTHREADS first so it picks up the right size, + # then call it here so later when it's called in libcudf it's a no-op. The explicit + # kvikio.defaults.set below handles subsequent calls (call_once only fires once). os.environ["KVIKIO_NTHREADS"] = str(nthreads) pylibcudf.io.kvikio.set_up_kvikio() kvikio.defaults.set( From 6506813ca35345e99c7d7aa002480107120b1fda Mon Sep 17 00:00:00 2001 From: Matthew Murray Date: Wed, 19 Aug 2026 12:14:13 +0000 Subject: [PATCH 5/5] dont init cuda ctx too early --- python/cudf_polars/cudf_polars/engine/dask.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/python/cudf_polars/cudf_polars/engine/dask.py b/python/cudf_polars/cudf_polars/engine/dask.py index 7552a1f63239..5ddb8efd22c9 100644 --- a/python/cudf_polars/cudf_polars/engine/dask.py +++ b/python/cudf_polars/cudf_polars/engine/dask.py @@ -374,7 +374,6 @@ def _setup_worker( """ assert dask_worker is not None - configure_kvikio(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]