Fix kvikio thread count override and the default backend to EASY_THREADPOOL in cudf-polars

docs/cudf/source/pylibcudf/api_docs/io/index.rst

@@ -18,6 +18,7 @@ I/O Functions
     avro
     csv
     json
+    kvikio
     orc
     parquet
     parquet_metadata

docs/cudf/source/pylibcudf/api_docs/io/kvikio.rst

@@ -0,0 +1,6 @@
+======
+KvikIO
+======
+
+.. automodule:: pylibcudf.io.kvikio
+   :members:

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

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,14 +374,14 @@ 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)
 
     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:

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:

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"

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"

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"

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,23 @@ 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 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(
+        {
+            "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 +761,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

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"

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
 )
 

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,

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

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()

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: ...

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()

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